#!/bin/bash

# NiuLink CentOS ISO UEFI Fix Script
# 修复UEFI安装时的fedora EFI目录错误
# 版本: v1.2
# 日期: 2025-06-06

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查参数
if [ $# -ne 1 ]; then
    echo "使用方法: $0 <original_iso_file>"
    echo "示例: $0 NiuLinkOS-original.iso"
    exit 1
fi

ORIGINAL_ISO="$1"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WORK_DIR="$SCRIPT_DIR/uefi_fix_work"
OUTPUT_ISO="${ORIGINAL_ISO%.*}-uefi-fixed.iso"

# 检查输入文件
if [ ! -f "$ORIGINAL_ISO" ]; then
    log_error "ISO文件不存在: $ORIGINAL_ISO"
    exit 1
fi

log_info "开始修复NiuLink ISO UEFI安装错误"
log_info "输入ISO: $ORIGINAL_ISO"
log_info "输出ISO: $OUTPUT_ISO"

# 检查必需工具
log_info "检查必需工具..."
for tool in xorriso unsquashfs mksquashfs; do
    if ! command -v $tool >/dev/null 2>&1; then
        log_error "缺少必需工具: $tool"
        log_info "请运行: sudo apt install -y xorriso squashfs-tools"
        exit 1
    fi
done
log_success "所有必需工具已安装"

# 清理旧的工作目录
if [ -d "$WORK_DIR" ]; then
    log_info "清理旧的工作目录..."
    sudo rm -rf "$WORK_DIR"
fi

# 创建工作目录
log_info "创建工作目录..."
mkdir -p "$WORK_DIR"
cd "$WORK_DIR"

# 提取ISO内容
log_info "提取ISO内容..."
xorriso -osirrox on -indev "$SCRIPT_DIR/$ORIGINAL_ISO" -extract / ./ >/dev/null 2>&1
log_success "ISO内容提取完成"

# 提取squashfs
log_info "提取squashfs文件系统..."
sudo unsquashfs -d rootfs LiveOS/squashfs.img >/dev/null 2>&1
log_success "squashfs提取完成"

# 挂载rootfs
log_info "挂载rootfs进行修改..."
ROOTFS_MOUNT=$(mktemp -d)
sudo mount -o loop rootfs/LiveOS/rootfs.img "$ROOTFS_MOUNT"
log_success "rootfs挂载完成: $ROOTFS_MOUNT"

# 备份原始文件
log_info "备份原始anaconda文件..."
ANACONDA_DIR="$ROOTFS_MOUNT/usr/lib64/python2.7/site-packages/pyanaconda/installclasses"

if [ ! -d "$ANACONDA_DIR" ]; then
    log_error "找不到anaconda安装类目录: $ANACONDA_DIR"
    sudo umount "$ROOTFS_MOUNT"
    exit 1
fi

sudo cp "$ANACONDA_DIR/fedora.py" "$ANACONDA_DIR/fedora.py.backup"
sudo cp "$ANACONDA_DIR/centos.py" "$ANACONDA_DIR/centos.py.backup"
log_success "原始文件备份完成"

# 应用核心修复
log_info "应用UEFI修复..."

# 1. 修复fedora.py的EFI目录
log_info "  -> 修复fedora.py EFI目录指向"
sudo sed -i 's/efi_dir = "fedora"/efi_dir = "centos"/' "$ANACONDA_DIR/fedora.py"

# 验证修复
if sudo grep -q 'efi_dir = "centos"' "$ANACONDA_DIR/fedora.py"; then
    log_success "     fedora.py EFI目录已修复为centos"
else
    log_error "     fedora.py EFI目录修复失败"
    sudo umount "$ROOTFS_MOUNT"
    exit 1
fi

# 2. 强制隐藏fedora安装类
log_info "  -> 强制隐藏fedora安装类"
sudo sed -i '/if productName\.startswith("Red Hat ") or productName\.startswith("CentOS"):/,+1c\    # Always hide fedora class to prevent EFI errors\n    hidden = True' "$ANACONDA_DIR/fedora.py"

# 验证隐藏
if sudo grep -q 'hidden = True' "$ANACONDA_DIR/fedora.py"; then
    log_success "     fedora安装类已强制隐藏"
else
    log_warning "     fedora安装类隐藏可能未完全生效"
fi

# 3. 确保CentOS安装类可见
log_info "  -> 确保CentOS安装类始终可见"
sudo sed -i 's/if not productName\.startswith("CentOS"):/if False:  # Always show CentOS class/' "$ANACONDA_DIR/centos.py"

# 验证CentOS类可见性
if sudo grep -q 'if False:.*Always show CentOS' "$ANACONDA_DIR/centos.py"; then
    log_success "     CentOS安装类已设置为始终可见"
else
    log_warning "     CentOS安装类可见性修改可能未完全生效"
fi

# 4. 确保系统标识正确
log_info "  -> 更新系统标识文件"
sudo tee "$ROOTFS_MOUNT/etc/os-release" > /dev/null << 'EOF'
NAME="CentOS Linux"
VERSION="7 (Core)"
ID="centos"
ID_LIKE="rhel fedora"
VERSION_ID="7"
PRETTY_NAME="CentOS Linux 7 (Core)"
ANSI_COLOR="0;31"
CPE_NAME="cpe:/o:centos:centos:7"
HOME_URL="https://www.centos.org/"
BUG_REPORT_URL="https://bugs.centos.org/"

CENTOS_MANTISBT_PROJECT="CentOS-7"
CENTOS_MANTISBT_PROJECT_VERSION="7"
REDHAT_SUPPORT_PRODUCT="centos"
REDHAT_SUPPORT_PRODUCT_VERSION="7"
EOF

log_success "UEFI修复应用完成"

# 卸载rootfs
log_info "卸载rootfs..."
sudo umount "$ROOTFS_MOUNT"
sudo rmdir "$ROOTFS_MOUNT"
log_success "rootfs卸载完成"

# 重新打包squashfs
log_info "重新打包squashfs..."
sudo mksquashfs rootfs LiveOS/squashfs.img -comp xz -noappend >/dev/null 2>&1
log_success "squashfs重新打包完成"

# 创建修复后的ISO
log_info "创建修复后的ISO..."
sudo xorriso -as mkisofs \
    -o "$SCRIPT_DIR/$OUTPUT_ISO" \
    -b isolinux/isolinux.bin \
    -c isolinux/boot.cat \
    -no-emul-boot \
    -boot-load-size 4 \
    -boot-info-table \
    -eltorito-alt-boot \
    -e images/efiboot.img \
    -no-emul-boot \
    -isohybrid-gpt-basdat \
    -V "ISOCDROM" \
    . >/dev/null 2>&1

if [ $? -eq 0 ]; then
    log_success "ISO创建完成: $OUTPUT_ISO"
else
    log_error "ISO创建失败"
    exit 1
fi

# 清理工作目录
log_info "清理工作目录..."
cd "$SCRIPT_DIR"
sudo rm -rf "$WORK_DIR"
log_success "清理完成"

# 显示结果
echo
echo "========================================"
echo "        UEFI修复完成"
echo "========================================"
echo "输入文件: $ORIGINAL_ISO"
echo "输出文件: $OUTPUT_ISO"
echo "文件大小: $(du -h "$OUTPUT_ISO" | cut -f1)"
echo
echo "修复内容:"
echo "✅ fedora.py EFI目录修复为centos"
echo "✅ fedora安装类强制隐藏"
echo "✅ CentOS安装类强制可见"
echo "✅ 系统标识确认为CentOS"
echo "✅ 保持UEFI + Legacy BIOS双启动"
echo
echo "预期效果:"
echo "• 不再出现 '/boot/efi/EFI/fedora/user.cfg' 错误"
echo "• UEFI安装过程将正常完成"
echo "• 支持Secure Boot环境"
echo
echo "使用方法:"
echo "1. 使用 $OUTPUT_ISO 创建安装介质"
echo "2. 支持UEFI和Legacy BIOS启动"
echo "3. 安装过程中选择CentOS安装选项"
echo
echo "建议下一步:"
echo "./verify_uefi_fix.sh \"$OUTPUT_ISO\""
echo "========================================"


